import moviepy.editor as mpy
import matplotlib.pyplot as plt, numpy as np, seaborn as sns, tqdm, pathlib, pims, cv2 as cv, xarray as xf, pandas as pd, scipy.ndimage as ndimage

vid = pims.Video('observations/4_data_settling_velocities/videos/17-04-12_20-Jun-19_settling_velocity.mov')
nframes = vid._len
start = 0
cent = np.zeros((2,nframes))
dt = 1/120
pix2mm = 200 / vid[0][:,120:-20].shape[1]

for i in tqdm.tqdm(range(start,start+nframes)):
    img = vid[i]
    img = img[:,600:850]
    img = cv.cvtColor(img, cv.COLOR_RGB2GRAY)
    img = cv.medianBlur(img, 7)
    img = cv.adaptiveThreshold(img, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 201, -41)
    img = cv.medianBlur(img, 5)
    kernel = np.ones((3,3),np.uint8)
    img = cv.morphologyEx(img, cv.MORPH_OPEN, kernel, iterations=3)
    img = cv.morphologyEx(img, cv.MORPH_CLOSE, kernel, iterations=3)
    ret, markers = cv.connectedComponents(img)
    max_shape = 0
    max_i = 0
    for k in range(1,markers.max()+1):
        if markers[markers==k].sum()/k > max_shape:
            max_shape = markers[markers==i].sum()/k
            max_i = k
    img = np.where(markers == max_i, 1, 0).astype(np.uint8)
    if max_i > 0:
        cent[:,i-start] = ndimage.measurements.center_of_mass(img)

    else:
        cent[:,i-start] = np.nan

y_pos = (cent[1,:]).astype(np.float) * pix2mm
x_pos = cent[0,:].astype(np.float) * pix2mm
x_vel = np.zeros_like(x_pos)
y_vel = np.zeros_like(y_pos)
x_vel[0] = (x_pos[1] - x_pos[0]) / dt
x_vel[1] = (x_pos[2] - x_pos[1]) / dt
x_vel[-2] = (x_pos[-2] - x_pos[-3]) / dt
x_vel[-1] = (x_pos[-1] - x_pos[-2]) / dt
y_vel[0] = (y_pos[1] - y_pos[0]) / dt
y_vel[1] = (y_pos[2] - y_pos[1]) / dt
y_vel[-2] = (y_pos[-2] - y_pos[-3]) / dt
y_vel[-1] = (y_pos[-1] - y_pos[-2]) / dt
for ii in tqdm.tqdm(range(2,nframes-2)):
    x_vel[ii] = (-x_pos[ii+2] + 8*x_pos[ii+1] - 8*x_pos[ii-1] + x_pos[ii-2]) / (12*dt)
    y_vel[ii] = (-y_pos[ii+2] + 8*y_pos[ii+1] - 8*y_pos[ii-1] + y_pos[ii-2]) / (12*dt)

# vel_mag = np.sqrt(x_vel**2 + y_vel**2)
vel_mag = np.sqrt(y_vel**2)
time = (np.ones_like(vel_mag) * dt).cumsum()
df = pd.DataFrame({'time': time, 'v': vel_mag})
df['v_smooth'] = df.v.interpolate(limit=11).rolling(window=31, min_periods=11, center=True).mean()
df['particleId'] = np.zeros_like(df.v.values)
ii = 0
particleId = 0
start_list = np.where(np.isfinite(df.v.values))[0]
while ii < df.shape[0]:
    if np.isfinite(df.v_smooth.values[ii]):
        df['particleId'][ii] = particleId
        ii += 1
    else:
        particleId += 1
        ilist = start_list[start_list > ii]
        if len(ilist) > 0:
            ii = ilist[0]
        else:
            break

count = 0
plt.figure(count)
dfs = []
for name, group in df.groupby('particleId'):
    if (group.v.std() < 40) & (group.v.mean() > 220):
        count += 1
        plt.plot(group.v.values, '-')
        # plt.title(group.v.std())
        dfs.append(group)
        plt.ylim(0, 1000)

v_mean = pd.concat(dfs).groupby('particleId').mean().v
v_std = pd.concat(dfs).groupby('particleId').std().v
print(v_mean.mean(), v_std.mean())
plt.title('mean v: %g +/- %g' % (v_mean.mean(), v_std.mean()))
plt.savefig('observations/4_data_settling_velocities/individual_grain_material_processing_codes/terminal_settling_tracks_oc.pdf')

df = pd.DataFrame(dict(vel = v_mean, vel_e = v_std))
df.to_csv('observations/4_data_settling_velocities/oc_settling_velocities.csv')


# fps = 60
# num_images = nframes
# duration = num_images/fps
#
# def make_frame(t):
#     img = vid[int(t*fps)]
#     img = img[:,400:850]
#     img = cv.cvtColor(img, cv.COLOR_RGB2GRAY)
#     img = cv.medianBlur(img, 7)
#     img = cv.adaptiveThreshold(img, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 201, -41)
#     img = cv.medianBlur(img, 5)
#     kernel = np.ones((3,3),np.uint8)
#     img = cv.morphologyEx(img, cv.MORPH_OPEN, kernel, iterations=5)
#     img = cv.morphologyEx(img, cv.MORPH_CLOSE, kernel, iterations=5)
#     ret, markers = cv.connectedComponents(img)
#     max_shape = 0
#     max_i = 0
#     for i in range(1,markers.max()+1):
#         if markers[markers==i].sum()/i > max_shape:
#             max_shape = markers[markers==i].sum()/i
#             max_i = i
#
#     img = np.where(markers == max_i, 255, 0).astype(np.uint8)
#     img = cv.cvtColor(img, cv.COLOR_GRAY2RGB)
#     return img
#
# animation = mpy.VideoClip(make_frame, duration=duration)
# animation.write_videofile('observations/4_data_settling_velocities/videos/17-04-12_20-Jun-19_settling_velocity_converted.mp4', fps=fps) # export as video
